import asyncio

from sql_graph.text2sql_graph import make_graph


async def execute_graph():
    """执行该自定义的图工作流"""
    async with make_graph() as graph:
        # 画图
        while True:
            user_input = input("请user输入查询语句：")
            if user_input.lower() in ['q', 'exit', 'quit']:
                print("bye, dialog end!!!")
                break
            else:
                async for event in graph.astream({"messages": [{"role": "user", "content": user_input}]},
                                                 stream_mode='values'):
                    event["messages"][-1].pretty_print()




if __name__ == '__main__':
    """将异步函数 execute_graph() 作为主协程运行，处理可能包含多个异步步骤的图执行流程。"""
    asyncio.run(execute_graph())